package com.zhxq.member.socket;

import com.zhxq.ai.repository.ChatHistoryRepository;
import jakarta.annotation.Resource;
import jakarta.websocket.OnMessage;
import jakarta.websocket.OnOpen;
import jakarta.websocket.Session;
import jakarta.websocket.server.PathParam;
import jakarta.websocket.server.ServerEndpoint;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;

@Slf4j
@Component
@EnableWebSocket
@ServerEndpoint("/member/ws/know/{chatId}")
public class WebSocketKnowServer {

    private static Map<String, Session> sessionMap = new HashMap<>();

    @Resource
    private ChatClient pdfChatClient;

    @Resource
    private ChatHistoryRepository mongoChatHistoryRepository;

    @OnOpen
    public void onOpen(Session session, @PathParam("chatId") String chatId) {
        log.info("WebSocketServer onOpen" + session.getId());
        sessionMap.put(chatId, session);
    }

    @OnMessage
    public void onMessage(Session session, String message) {
        log.info("接收到了客户端 {} 发来的消息 : {}", session.getId(), message);
    }

    /**
     * 发送消息给AI
     * @param prompt
     * @param chatId
     */
    public void sendMessageToAI(String prompt, String chatId) {
        mongoChatHistoryRepository.save("knowledge", chatId);
        Flux<String> responseStream = pdfChatClient.prompt()
                .user(prompt)
                .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId))
                .stream()
                .content();
        Session session = sessionMap.get(chatId);
        responseStream.subscribe(
                data -> {
                    try {
                        session.getBasicRemote().sendText(data);
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                },
                err -> {
                    try {
                        session.getBasicRemote().sendText("[ERROR]" + err.getMessage());
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                    try {
                        session.close();
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                },
                () -> {
                    try {
                        session.getBasicRemote().sendText("[END]"); // 发送结束标识
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                }
        );
    }
}
